#!/usr/bin/env python3
from collections import deque
from dataclasses import dataclass
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from PyQt5.QtWidgets import (
    QApplication,
    QDialog,
    QTabWidget,
    QVBoxLayout,
    QWidget,
    QLabel,
    QProgressBar,
    QHBoxLayout,
    QGridLayout,
)
import os.path
import psutil
import sys
import time
import typing as tp
import pyqtgraph as pg


N_CPU_COLUMNS = 4
MAX_IFACE_LOAD = 1 * 1024 * 1024
MAX_DISK_LOAD = 1 * 1024 * 1024
MIN_SYMBOL_WIDTH = 10
PLOT_SIZE_S = 60
TIME_RANGE = list(range(-PLOT_SIZE_S, 0))
PLOT_PEN = pg.mkPen(color='b', width=2)


@dataclass
class HwItem:
    name: str
    update_callback: tp.Callable


def to_mb(bytes: tp.Any) -> int:
    return int(bytes / 1024 / 1024)


class CpuInfo(HwItem):
    @dataclass
    class Counters:
        core: int
        load: int

    def __init__(self):
        super().__init__("CPU", self.update)
        self.cpu_count = psutil.cpu_count(logical=True)
        self.update()

    def update(self):
        self.counters = []
        for i, percentage in enumerate(psutil.cpu_percent(percpu=True)):
            self.counters.append(self.Counters(i, int(percentage)))


class MemInfo(HwItem):
    @dataclass
    class Counters:
        device: str
        total: int
        used: int

    def __init__(self):
        super().__init__("Memory", self.update)
        self.update()

    def update(self):
        self.counters = []
        virt = psutil.virtual_memory()
        self.counters.append(self.Counters("mem", to_mb(virt.total), to_mb(virt.used)))
        swap = psutil.swap_memory()
        if swap.total > 0:
            self.counters.append(self.Counters("swap", to_mb(swap.total), to_mb(swap.used)))


class DiskInfo(HwItem):
    @dataclass
    class Counters:
        device: str
        size: int
        used: int
        rbytes: int
        wbytes: int
        rspeed: int = 0
        wspeed: int = 0

    def __init__(self):
        super().__init__("Disk", self.update)
        self.update_time = time.monotonic()
        self.prev_counters: tp.Dict(str, self.Counters) = {}
        self.update()

    def update(self):
        self.counters = []
        partitions = psutil.disk_partitions()
        io_counters = psutil.disk_io_counters(perdisk=True)
        curr_time = time.monotonic()
        delta_time = curr_time - self.update_time

        check_devices = []
        for partition in partitions:
            if "/dev/loop" in partition.device:
                continue
            if partition.device in check_devices:
                continue
            check_devices.append(partition.device)
            disk_usage = psutil.disk_usage(partition.mountpoint)

            counters = self.Counters(
                partition.device,
                to_mb(disk_usage.total),
                to_mb(disk_usage.used),
                io_counters[os.path.basename(partition.device)].read_bytes,
                io_counters[os.path.basename(partition.device)].write_bytes,
            )

            self.counters.append(counters)
            if counters.device in self.prev_counters:
                prev_counters = self.prev_counters[counters.device]
                counters.rspeed = int((counters.rbytes - prev_counters.rbytes) / delta_time)
                counters.wspeed = int((counters.wbytes - prev_counters.wbytes) / delta_time)
            self.prev_counters[counters.device] = counters

        self.update_time = curr_time


class NetworkInfo(HwItem):
    @dataclass
    class Counters:
        iface: str
        totalin: int
        totalout: int
        speedin: int = 0
        speedout: int = 0

    def __init__(self):
        super().__init__("Disk", self.update)
        self.update_time = time.monotonic()
        self.prev_counters: tp.Dict(str, self.Counters) = {}
        self.update()

    def update(self):
        self.counters = []
        netstat = psutil.net_io_counters(pernic=True)
        curr_time = time.monotonic()
        delta_time = curr_time - self.update_time
        for iface, stat in netstat.items():
            counters = self.Counters(iface, stat.bytes_sent, stat.bytes_recv)
            self.counters.append(counters)
            if counters.iface in self.prev_counters:
                prev_counters = self.prev_counters[counters.iface]
                counters.speedin = int((counters.totalin - prev_counters.totalin) / delta_time)
                counters.speedout = int((counters.totalout - prev_counters.totalout) / delta_time)
            self.prev_counters[counters.iface] = counters
        self.update_time = curr_time


class TabCpu(QWidget):
    def __init__(self, cpuinfo):
        super().__init__()
        self.cpuinfo = cpuinfo

        self.widgets = {}
        self.plotvalues = {}
        self.plots_and_labels = []
        for i, item in enumerate(self.cpuinfo.counters):
            hbox = QHBoxLayout()
            lab = QLabel(f"#{item.core:02d}:")
            hbox.addWidget(lab)

            plot = pg.PlotWidget()
            plot.setBackground("w")
            plot.setMouseEnabled(False, False)
            plot.setYRange(0, 100)
            plot.setXRange(-PLOT_SIZE_S, 0, padding=0)
            plot.getPlotItem().addLegend()
            hbox.addWidget(plot)

            self.widgets[item.core] = (plot,)
            self.plotvalues[item.core] = deque([0] * PLOT_SIZE_S, maxlen=PLOT_SIZE_S)
            self.plots_and_labels.append(hbox)

        grid = QGridLayout()
        for i, item in enumerate(self.plots_and_labels):
            grid.addLayout(item, int(i / N_CPU_COLUMNS), int(i % N_CPU_COLUMNS), Qt.AlignHCenter)
        self.setLayout(grid)

    def updateTab(self, unused: int):
        for item in self.cpuinfo.counters:
            plotvalues = self.plotvalues[item.core]
            plotvalues.rotate(-1)
            plotvalues[-1] = item.load
            self.widgets[item.core][0].plot(
                TIME_RANGE, plotvalues, clear=True, pen=PLOT_PEN, name=f"load: {item.load} %"
            )


class TabMemory(QWidget):
    def __init__(self, meminfo):
        super().__init__()
        self.meminfo = meminfo

        self.progressbars = {}
        vbox = QVBoxLayout()
        for val in self.meminfo.counters:
            hbox = QHBoxLayout()

            lab = QLabel(f"{val.device}")
            hbox.addWidget(lab)

            bar = QProgressBar()
            bar.setMaximum(val.total)

            bar.setValue(val.used)
            bar.setFormat("%v MB / %m MB")
            hbox.addWidget(bar)
            self.progressbars[val.device] = bar

            vbox.addLayout(hbox)
        self.setLayout(vbox)

    def updateTab(self, unused: int):
        for item in self.meminfo.counters:
            self.progressbars[item.device].setValue(item.used)


class TabDisk(QWidget):
    def __init__(self, diskinfo):
        super().__init__()
        self.diskinfo = diskinfo

        self.progressbars = {}
        self.rplotvalues = {}
        self.wplotvalues = {}
        vbox = QVBoxLayout()
        max_name_len = max([len(val.device) for val in self.diskinfo.counters])
        for val in self.diskinfo.counters:
            hbox = QHBoxLayout()
            lab = QLabel(f"{val.device}")
            lab.setFixedWidth(max_name_len * MIN_SYMBOL_WIDTH)
            hbox.addWidget(lab)

            bar = QProgressBar()
            bar.setMaximum(val.size)
            bar.setValue(val.used)
            bar.setFormat("%v MB / %m MB")
            hbox.addWidget(bar)

            rplot = pg.PlotWidget()
            rplot.setBackground("w")
            rplot.setMouseEnabled(False, False)
            rplot.enableAutoRange(pg.ViewBox.XAxis)
            rplot.setXRange(-PLOT_SIZE_S, 0, padding=0)
            rplot.getPlotItem().addLegend()
            hbox.addWidget(rplot)

            wplot = pg.PlotWidget()
            wplot.setBackground("w")
            wplot.setMouseEnabled(False, False)
            wplot.enableAutoRange(pg.ViewBox.XAxis)
            wplot.setXRange(-PLOT_SIZE_S, 0, padding=0)
            wplot.getPlotItem().addLegend()
            hbox.addWidget(wplot)

            self.progressbars[val.device] = (bar, rplot, wplot)

            vbox.addLayout(hbox)

            self.rplotvalues[val.device] = deque([0] * PLOT_SIZE_S, maxlen=PLOT_SIZE_S)
            self.wplotvalues[val.device] = deque([0] * PLOT_SIZE_S, maxlen=PLOT_SIZE_S)
        self.setLayout(vbox)

    def updateTab(self, unused: int):
        for item in self.diskinfo.counters:
            self.progressbars[item.device][0].setValue(item.used)
            rplotvalues = self.rplotvalues[item.device]
            rplotvalues.rotate(-1)
            rplotvalues[-1] = item.rspeed
            self.progressbars[item.device][1].plot(
                TIME_RANGE, rplotvalues, clear=True, pen=PLOT_PEN, name=f"Read speed, B/s: {item.rspeed}"
            )
            wplotvalues = self.wplotvalues[item.device]
            wplotvalues.rotate(-1)
            wplotvalues[-1] = item.wspeed
            self.progressbars[item.device][2].plot(
                TIME_RANGE, wplotvalues, clear=True, pen=PLOT_PEN, name=f"Write speed, B/s: {item.wspeed}"
            )


class TabNetwork(QWidget):
    def __init__(self, netwinfo):
        super().__init__()
        self.netwinfo = netwinfo

        self.progressbars = {}
        self.rplotvalues = {}
        self.wplotvalues = {}
        vbox = QVBoxLayout()
        max_name_len = max([len(val.iface) for val in self.netwinfo.counters])
        for val in self.netwinfo.counters:
            hbox = QHBoxLayout()
            lab = QLabel(f"{val.iface}")
            lab.setFixedWidth(max_name_len * MIN_SYMBOL_WIDTH)
            hbox.addWidget(lab)

            rplot = pg.PlotWidget()
            rplot.setBackground("w")
            rplot.setMouseEnabled(False, False)
            rplot.enableAutoRange(pg.ViewBox.XAxis)
            rplot.setXRange(-PLOT_SIZE_S, 0, padding=0)
            rplot.getPlotItem().addLegend()
            hbox.addWidget(rplot)

            wplot = pg.PlotWidget()
            wplot.setBackground("w")
            wplot.setMouseEnabled(False, False)
            wplot.enableAutoRange(pg.ViewBox.XAxis)
            wplot.setXRange(-PLOT_SIZE_S, 0, padding=0)
            wplot.getPlotItem().addLegend()
            hbox.addWidget(wplot)
            self.progressbars[val.iface] = (rplot, wplot)

            vbox.addLayout(hbox)
            self.rplotvalues[val.iface] = deque([0] * PLOT_SIZE_S, maxlen=PLOT_SIZE_S)
            self.wplotvalues[val.iface] = deque([0] * PLOT_SIZE_S, maxlen=PLOT_SIZE_S)
        self.setLayout(vbox)

    def updateTab(self, unused: int):
        for item in self.netwinfo.counters:
            rplotvalues = self.rplotvalues[item.iface]
            rplotvalues.rotate(-1)
            rplotvalues[-1] = item.speedin
            self.progressbars[item.iface][0].plot(
                TIME_RANGE, rplotvalues, clear=True, pen=PLOT_PEN, name=f"Input flow, B/s: {item.speedin}"
            )
            wplotvalues = self.wplotvalues[item.iface]
            wplotvalues.rotate(-1)
            wplotvalues[-1] = item.speedout
            self.progressbars[item.iface][1].plot(
                TIME_RANGE, wplotvalues, clear=True, pen=PLOT_PEN, name=f"Output flow, B/s: {item.speedout}"
            )


class MainForm(QDialog):
    cpuinfo = CpuInfo()
    meminfo = MemInfo()
    diskinfo = DiskInfo()
    netwinfo = NetworkInfo()

    class CollectorThread(QThread):
        change_value = pyqtSignal(int)

        def run(self):
            next_update_time = time.monotonic() + 1
            while True:
                MainForm.cpuinfo.update()
                MainForm.meminfo.update()
                MainForm.diskinfo.update()
                MainForm.netwinfo.update()
                self.change_value.emit(1)

                curr_time = time.monotonic()
                sleep_time = next_update_time - curr_time
                time.sleep(sleep_time if sleep_time > 0 else 0)
                next_update_time += 1

    def __init__(self):
        super().__init__()
        self.setWindowTitle("Hardware monitor")

        tab_cpu = TabCpu(MainForm.cpuinfo)
        tab_memory = TabMemory(MainForm.meminfo)
        tab_disk = TabDisk(MainForm.diskinfo)
        tab_network = TabNetwork(MainForm.netwinfo)

        vbox = QVBoxLayout()
        tabWidget = QTabWidget()
        tabWidget.addTab(tab_cpu, "CPU")
        tabWidget.addTab(tab_memory, "Memory")
        tabWidget.addTab(tab_disk, "Disk")
        tabWidget.addTab(tab_network, "Network")

        self.thread = self.CollectorThread()
        self.thread.change_value.connect(tab_cpu.updateTab)
        self.thread.change_value.connect(tab_memory.updateTab)
        self.thread.change_value.connect(tab_disk.updateTab)
        self.thread.change_value.connect(tab_network.updateTab)
        self.thread.start()

        vbox.addWidget(tabWidget)
        self.setLayout(vbox)


if __name__ == "__main__":
    app = QApplication(sys.argv)
    tabdialog = MainForm()
    tabdialog.show()
    app.exec_()
